from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim
from configs import M3E_BASE_MODEL_PATH, device
from faiss import normalize_L2
from tqdm import trange
import numpy as np
import torch
import logging
from torch import Tensor
# M3E_BASE_MODEL_PATH = "../embedding_models/m3e-base"
# device = torch.device("cuda:3")
logger = logging.getLogger(__name__)

show_progress_bar = (logger.getEffectiveLevel()==logging.INFO or logger.getEffectiveLevel()==logging.DEBUG)
model = SentenceTransformer(M3E_BASE_MODEL_PATH)
model.eval()
model = model.to(device)


def batch_to_device(batch):
    for key in batch:
        if isinstance(batch[key], Tensor):
            batch[key] = batch[key].to(device)
    return batch


def m3e_encode(text_list, is_normalize=True):
    batch_size = 32
    if not isinstance(text_list, list):
        text_list = [text_list]
        flag = 0
    else:
        flag = 1
    length_sorted_idx = np.argsort([-len(sen) for sen in text_list])
    sentences_sorted = [text_list[idx] for idx in length_sorted_idx]
    all_embeddings = []
    with torch.no_grad():
        for start_index in trange(0, len(text_list), batch_size, desc="Batches", disable=not show_progress_bar):
            sentences_batch = sentences_sorted[start_index:start_index+batch_size]
            encoded_input = model.tokenize(sentences_batch)
            encoded_input = batch_to_device(encoded_input)
            fea = model(encoded_input)
            embeddings = fea["sentence_embedding"]
            embeddings = embeddings.detach()
            if is_normalize:
                embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
            embeddings = embeddings.cpu()
            all_embeddings.extend(embeddings)
    all_embeddings = [all_embeddings[idx] for idx in np.argsort(length_sorted_idx)]
    all_embeddings = np.asarray([emb.numpy() for emb in all_embeddings])
    if flag == 0:
        return all_embeddings[0]
    else:
        return all_embeddings
    return all_embeddings

# device = torch.device("cuda:3")
# m3e_base_model = SentenceTransformer(M3E_BASE_MODEL_PATH).to(device)

# if __name__ == '__main__':
#     sentences = [
#         '* Moka 此文本嵌入模型由 MokaAI 训练并开源，训练脚本使用 uniem',
#         '* Massive 此文本嵌入模型通过**千万级**的中文句对数据集进行训练',
#         '* Mixed 此文本嵌入模型支持中英双语的同质文本相似度计算，异质文本检索等功能，未来还会支持代码检索，ALL in one'
#     ]
#     embeddings = m3e_base_model.encode(sentences)
#     print(embeddings.shape)
#     for sentence, embedding in zip(sentences, embeddings):
#         print("Sentence:", sentence)
#         print("Embedding:", embedding.shape)
#         print("")
#     print(np.linalg.norm(embeddings, axis=1, keepdims=True).shape)
#     embeddings_norm = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
#     print(embeddings_norm)

    

# sentences = "大都会人寿恶性肿瘤短期疾病保险 保险金申请 ', 'answer': '一、恶性肿瘤保险金的申请 恶性肿瘤保险金的申请人为恶性肿瘤保险金受益人。 在申请恶性肿瘤保险金时，申请人须填写保险金给付申请书，并提供下列证明和资料： （1）保险合同； （2）申请人的有效身份证件； （3）国家卫生行政部门认定的医疗机构出具的被保险人病理组织学检查、血液检查及其他科学方法检验报告的疾病诊断证明书； （4）申请人与被保险人的相关关系证明（如有需要）。 二、其他 上述相关证明和资料，除保险合同外，我们审核原件，审核完毕后留存复印件，原件返还给申请人或受托人。 保险金作为被保险人遗产时，必须提供可证明合法继承权的相关权利文件。 以上证明和资料不完整的，我们将及时一次性通知申请人补充提供有关的证明和资料。 除有关法律、行政法规不允许外，我们将保留进行医学鉴定或核实的权利。"
# print(m3e_encode(sentences).dtype)
# print(m3e_encode(sentences).shape)
# embed = m3e_encode("被保险人被宣告死亡时如何处理?")
# print(embed[:10])
# print(model.encode("被保险人被宣告死亡时如何处理?")[:10])
# np.save("query1.npy", embed)
# import json
# with open("../../faiss.json", "r", encoding="utf-8") as f:
#     q = json.load(f)
# q_list = []
# for _, v in q.items():
#     question = v["question"]
#     if len(question.encode("utf-8")) <= 2048:
#         q_list.append(question)
# embeddings = m3e_encode(q_list)
# np.save("question.npy", embeddings)
# with open("question.txt", "w", encoding="utf-8") as f:
#     for q in q_list:
#         f.write(q)
#         f.write("\n")
